from typing import List, Dict, Any, Optional


class ConversationMemory:
    """对话内存管理类"""
    
    def __init__(self, max_history: int = 50):
        """
        初始化对话内存
        
        Args:
            max_history: 最大保存的对话历史数量
        """
        self.max_history = max_history
        self.conversation_history: List[Dict[str, Any]] = []
        self.session_id: Optional[str] = None
    
    def add_message(self, role: str, content: str, metadata: Optional[Dict] = None):
        """
        添加消息到对话历史
        
        Args:
            role: 角色 (user/assistant)
            content: 消息内容
            metadata: 额外的元数据
        """
        message = {
            "role": role,
            "content": content,
            "timestamp": self._get_current_timestamp(),
            "metadata": metadata or {}
        }
        
        self.conversation_history.append(message)
        
        # 保持历史记录在限制范围内
        if len(self.conversation_history) > self.max_history:
            self.conversation_history = self.conversation_history[-self.max_history:]
    
    def get_conversation_history(self, last_n: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        获取对话历史
        
        Args:
            last_n: 获取最近的n条消息，None表示获取全部
            
        Returns:
            对话历史列表
        """
        if last_n is None:
            return self.conversation_history.copy()
        
        return self.conversation_history[-last_n:] if last_n > 0 else []
    
    def get_formatted_history(self, last_n: Optional[int] = None) -> str:
        """
        获取格式化的对话历史字符串
        
        Args:
            last_n: 获取最近的n条消息
            
        Returns:
            格式化的对话历史字符串
        """
        history = self.get_conversation_history(last_n)
        
        if not history:
            return "暂无对话历史"
        
        formatted_messages = []
        for msg in history:
            role_name = "用户" if msg["role"] == "user" else "助手"
            formatted_messages.append(f"{role_name}: {msg['content']}")
        
        return "\n".join(formatted_messages)
    
    def clear_history(self):
        """清空对话历史"""
        self.conversation_history.clear()
    
    def get_last_user_message(self) -> Optional[str]:
        """获取最后一条用户消息"""
        for msg in reversed(self.conversation_history):
            if msg["role"] == "user":
                return msg["content"]
        return None
    
    def get_last_assistant_message(self) -> Optional[str]:
        """获取最后一条助手消息"""
        for msg in reversed(self.conversation_history):
            if msg["role"] == "assistant":
                return msg["content"]
        return None
    
    def get_conversation_summary(self) -> Dict[str, Any]:
        """获取对话摘要信息"""
        total_messages = len(self.conversation_history)
        user_messages = len([msg for msg in self.conversation_history if msg["role"] == "user"])
        assistant_messages = len([msg for msg in self.conversation_history if msg["role"] == "assistant"])
        
        return {
            "total_messages": total_messages,
            "user_messages": user_messages,
            "assistant_messages": assistant_messages,
            "session_id": self.session_id,
            "last_message_time": self.conversation_history[-1]["timestamp"] if self.conversation_history else None
        }
    
    def set_session_id(self, session_id: str):
        """设置会话ID"""
        self.session_id = session_id
    
    def _get_current_timestamp(self) -> str:
        """获取当前时间戳"""
        import datetime
        return datetime.datetime.now().isoformat()


class ContextManager:
    """上下文管理器"""
    
    def __init__(self):
        self.context_data: Dict[str, Any] = {}
        self.user_preferences: Dict[str, Any] = {}
    
    def set_context(self, key: str, value: Any):
        """设置上下文数据"""
        self.context_data[key] = value
    
    def get_context(self, key: str, default: Any = None) -> Any:
        """获取上下文数据"""
        return self.context_data.get(key, default)
    
    def update_context(self, updates: Dict[str, Any]):
        """批量更新上下文"""
        self.context_data.update(updates)
    
    def clear_context(self):
        """清空上下文"""
        self.context_data.clear()
    
    def set_user_preference(self, key: str, value: Any):
        """设置用户偏好"""
        self.user_preferences[key] = value
    
    def get_user_preference(self, key: str, default: Any = None) -> Any:
        """获取用户偏好"""
        return self.user_preferences.get(key, default)
    
    def get_full_context(self) -> Dict[str, Any]:
        """获取完整上下文信息"""
        return {
            "context_data": self.context_data.copy(),
            "user_preferences": self.user_preferences.copy()
        }
